#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/10/8 14:35
# @Author  : Crissu
# @Site    : 
# @File    : hsv.py
# @Software: PyCharm

import cv2, os, multiprocessing
import numpy as np
from configs import hsv_config
from utils import universal_util

def ImageHSV(mode='S'):
    # 模式选择
    if mode == 'S':
        dealFunc = doImageHSV_S
    elif mode == 'V':
        dealFunc = doImageHSV_V
    elif mode == 'H':
        dealFunc = doImageHSV_H

    classList = os.listdir(hsv_config.FromPath)
    print(classList)
    universal_util.MakeFloder(hsv_config.ToPath)
    pool = multiprocessing.Pool(16)
    pool.map(dealFunc, classList)

'''
定义hsv变换函数：
hue_delta是色调变化比例
sat_delta是饱和度变化比例
val_delta是明度变化比例
'''
def hsv_transform(img, hue_delta, sat_mult, val_mult):
    img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float)
    img_hsv[:, :, 0] = (img_hsv[:, :, 0] + hue_delta) % 180
    img_hsv[:, :, 1] *= sat_mult
    img_hsv[:, :, 2] *= val_mult
    img_hsv[img_hsv > 255] = 255
    return cv2.cvtColor(np.round(img_hsv).astype(np.uint8), cv2.COLOR_HSV2BGR)

'''
随机hsv变换
hue_vari是色调变化比例的范围
sat_vari是饱和度变化比例的范围
val_vari是明度变化比例的范围
'''
def random_hsv_transform(img, hue_vari, sat_vari, val_vari):
    hue_delta = np.random.randint(-hue_vari, hue_vari)
    sat_mult = 1 + np.random.uniform(sat_vari, sat_vari)
    val_mult = 1 + np.random.uniform(val_vari, val_vari)
    return hsv_transform(img, hue_delta, sat_mult, val_mult)

def doImageHSV_V(className):
    valMap = {
              "120": [0.2, 0.4, 0.6, 0.8, 1],
              "150": [-0.1, -0.2, 0.2, 0.4, 0.6],
              "255": [-0.1, -0.2, -0.3, -0.4, -0.5],
              }
    path = hsv_config.FromPath + className + "/"
    nameList = os.listdir(path)
    for name in nameList:
        rawPath = path + name
        img = cv2.imread(rawPath)
        # val_mult = [-0.2, 0.2, 0.4, 0.6, 0.8, 1] #, 1.5, 2, 2.5, 3
        img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float)
        array = np.array(img_hsv[:, :, 2])
        val = int(np.mean(np.mean(array, axis=1), axis=0))
        # print(val)
        # if val > 0 and val <= 30:
        #     val_mult = valMap["30"]
        # elif val > 30 and val <= 60:
        #     val_mult = valMap["60"]
        # elif val > 60 and val <= 90:
        #     val_mult = valMap["90"]
        if val <= 120:
            val_mult = valMap["120"]
        elif val > 120 and val <= 150:
            val_mult = valMap["150"]
        elif val > 150 and val <= 255:
            val_mult = valMap["255"]

        for i in val_mult:
            # img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float)
            # print(img_hsv[:, :, 2])
            # array = np.array(img_hsv[:, :, 2])
            # print(int(np.mean(np.mean(array, axis=1), axis=0)))
            saveImg = random_hsv_transform(img, 1, 0.1, i)
            saveName = "v_" + str(i) + name
            savePath = hsv_config.ToPath + className + "/" + saveName
            cv2.imwrite(savePath, saveImg)
            print("doImageHSV_V - ", savePath, " - done")
            # img_hsv = cv2.cvtColor(saveImg, cv2.COLOR_BGR2HSV).astype(np.float)
            # print(img_hsv[:, :, 2])
            # array = np.array(img_hsv[:, :, 2])
            # print(int(np.mean(np.mean(array, axis=1), axis=0)))


def doImageHSV_S(className):
    sat_mult = [-0.4, -0.2, 0.4, 0.6, 0.8]
    path = hsv_config.FromPath + className + "/"
    nameList = os.listdir(path)
    nameLen = len(os.listdir(hsv_config.SourcePath + className + '/'))
    if nameLen >150:
        sat_mult = [-0.4, -0.2, 0.4, 0.6]
    for name in nameList:
        rawPath = path + name
        img = cv2.imread(rawPath)
        # img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float)
        for i in sat_mult:
            # array = np.array(img_hsv[:, :, 1])
            # print(int(np.mean(np.mean(array, axis=1), axis=0)))
            saveImg = random_hsv_transform(img, 1, i, 0)
            saveName = "s_" + str(i) + name
            savePath = hsv_config.ToPath + className + "/" + saveName
            cv2.imwrite(savePath, saveImg)
            print("doImageHSV_S - ", savePath, " - done")
            # img_hsv = cv2.cvtColor(saveImg, cv2.COLOR_BGR2HSV).astype(np.float)
            # array = np.array(img_hsv[:, :, 1])
            # print(int(np.mean(np.mean(array, axis=1), axis=0)))

def doImageHSV_H(className):
    hue_delta = [2, 4, 6, 8, 9]
    path = hsv_config.FromPath + className + "/"
    nameList = os.listdir(path)
    nameLen = len(os.listdir(hsv_config.SourcePath + className + '/'))
    if nameLen > 100:
        return
    if nameLen>50 and nameLen<=100:
        hue_delta = [4, 8]
    elif nameLen>40 and nameLen<=50:
        hue_delta = [2, 4, 6, 8]
    elif nameLen>=20 and nameLen<=40:
        hue_delta = [2, 4, 6, 8, 9]
    for name in nameList:
        rawPath = path + name
        img = cv2.imread(rawPath)
        # img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float)
        for i in hue_delta:
            # array = np.array(img_hsv[:, :, 0])
            # print(int(np.mean(np.mean(array, axis=1), axis=0)))
            saveImg = random_hsv_transform(img, i, 0.1, 0)
            saveName = "h_" + str(i) + name
            savePath = hsv_config.ToPath + className + "/" + saveName
            cv2.imwrite(savePath, saveImg)
            print("doImageHSV_H - ", savePath, " - done")
            # img_hsv = cv2.cvtColor(saveImg, cv2.COLOR_BGR2HSV).astype(np.float)
            # array = np.array(img_hsv[:, :, 0])
            # print(int(np.mean(np.mean(array, axis=1), axis=0)))

# if __name__ == '__main__':
    # valMap = {"30": [0.5, 1, 1.5, 2, 2.5, 3],
    #           "60": [0.5, 1, 1.5, 2, 2.5],
    #           "90": [-1, 0.5, 1, 1.5, 2],
    #           "120": [0.2, 0.4, 0.6, 0.8, 1],
    #           "150": [-0.2, 0.2, 0.4, 0.6],
    #           "255": [-0.1, -0.2, -0.3, -0.4, -0.5],
    #           }
    # name = "X3_MWSnap970.jpg"
    # img = cv2.imread("test/"+name)
    # val_mult = [2,4,6,8,9]
    # # val_mult = [-1, -2, -3, -4, -5]
    # for i in val_mult:
    #     saveImg = random_hsv_transform(img, i, 0.1, 0)
    #     saveName = "s_" + str(i) + name
    #     cv2.imwrite("test/"+saveName, saveImg)


    # path = "C:/sjj/workspace/dataset/linyehaichong/croped_selected_resized_flip_rotate_hsv/77tou_ming_shu_guang_la_chan/"
    # nameList = os.listdir(path)
    # for name in nameList:
    #     img = cv2.imread(path+name)
    #     # val_mult = [-0.2, 0.2, 0.4, 0.6, 0.8, 1] #, 1.5, 2, 2.5, 3
    #     img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float)
    #     array = np.array(img_hsv[:, :, 2])
    #     val = int(np.mean(np.mean(array, axis=1), axis=0))
    #     if val >0 and val <=30:
    #         val_mult = valMap["30"]
    #     elif val >30 and val<=60:
    #         val_mult = valMap["60"]
    #     elif val >60 and val<=90:
    #         val_mult = valMap["90"]
    #     elif val >90 and val<=120:
    #         val_mult = valMap["120"]
    #     elif val >120 and val<=150:
    #         val_mult = valMap["150"]
    #     elif val >150 and val<=255:
    #         val_mult = valMap["255"]
    #
    #     for i in val_mult:
    #         img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float)
    #         print(img_hsv[:, :, 2])
    #         array = np.array(img_hsv[:, :, 2])
    #         print(int(np.mean(np.mean(array, axis=1), axis=0)))
    #         saveImg = random_hsv_transform(img, 1, 0.1, i)
    #         saveName = str(i) + name
    #         cv2.imwrite("test/"+saveName, saveImg)
    #         img_hsv = cv2.cvtColor(saveImg, cv2.COLOR_BGR2HSV).astype(np.float)
    #         print(img_hsv[:, :, 2])
    #         array = np.array(img_hsv[:, :, 2])
    #         print(int(np.mean(np.mean(array, axis=1), axis=0)))

    # name = "0.4MWSnap155.jpg"
    # path = "test/" + name
    # img = cv2.imread(path)
    # sat_mult = [-0.4, -0.2, 0.4, 0.8, 1.2] #, 1.5, 2, 2.5, 3
    # for i in sat_mult:
    #     img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float)
    #     print(img_hsv[:, :, 1])
    #     array = np.array(img_hsv[:, :, 1])
    #     print(int(np.mean(np.mean(array, axis=1), axis=0)))
    #     saveImg = random_hsv_transform(img, 1, i, 0)
    #     saveName = str(i) + name
    #     cv2.imwrite("test/"+saveName, saveImg)
    #     img_hsv = cv2.cvtColor(saveImg, cv2.COLOR_BGR2HSV).astype(np.float)
    #     print(img_hsv[:, :, 1])
    #     array = np.array(img_hsv[:, :, 1])
    #     print(int(np.mean(np.mean(array, axis=1), axis=0)))










